#include <bsl_err_internal.h>
#include <bsl_errno.h>
#include <bsl_params.h>
#include <bsl_sal.h>
#include <crypt_algid.h>
#include <crypt_errno.h>
#include <crypt_rsa.h>
#include <crypt_utils.h>
#include <securec.h>

#include "kaev1_rsa.h"
#include "kaev1_rsa_local.h"
#include "kaep_err.h"

int32_t RSA_CheckBitsLength(uint32_t bits)
{
    switch (bits) {
        case 1024:
        case 2048:
        case 3072:
        case 4096:
            return CRYPT_SUCCESS;
        default:
            return KAE_WD_ERR_KEY_BITS;
    }
}

int32_t GetRsaParam(const BSL_Param *params, int32_t type, const uint8_t **value, uint32_t *valueLen)
{
    const BSL_Param *temp = BSL_PARAM_FindConstParam(params, type);
    if (temp == NULL || temp->valueLen == 0 || temp->value == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_INVALID_ARG);
        return CRYPT_INVALID_ARG;
    }

    *value = temp->value;
    *valueLen = temp->valueLen;
    return CRYPT_SUCCESS;
}

int32_t GetRsaBits(const BSL_Param *params, uint32_t *bits)
{
    uint32_t bitsLen = sizeof(*bits);
    const BSL_Param *temp = BSL_PARAM_FindConstParam(params, CRYPT_PARAM_RSA_BITS);
    if (temp == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_INVALID_ARG);
        return CRYPT_INVALID_ARG;
    }

    int32_t ret = BSL_PARAM_GetValue(temp, CRYPT_PARAM_RSA_BITS, BSL_PARAM_TYPE_UINT32, bits, &bitsLen);
    if (ret != BSL_SUCCESS || *bits < RSA_MIN_MODULUS_BITS || *bits > RSA_MAX_MODULUS_BITS) {
        BSL_ERR_PUSH_ERROR(CRYPT_INVALID_ARG);
        return CRYPT_INVALID_ARG;
    }

    return CRYPT_SUCCESS;
}

int32_t ValidateRsaParams(uint32_t eLen, uint32_t bits)
{
    /* the length of e cannot be greater than bits */
    if (eLen > BN_BITS_TO_BYTES(bits)) {
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_KEY_BITS);
        return CRYPT_RSA_ERR_KEY_BITS;
    }
    return CRYPT_SUCCESS;
}

int32_t IsRSASetParaVaild(const KAE_RSA_Ctx *ctx, const RSA_Para *para)
{
    if (ctx == NULL || para == NULL || para->e == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    if (para->bits > RSA_MAX_MODULUS_BITS || para->bits < RSA_MIN_MODULUS_BITS) {
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_KEY_BITS);
        return CRYPT_RSA_ERR_KEY_BITS;
    }

    if (BN_GetBit(para->e, 0) != true || BN_IsLimb(para->e, 1) == true) {
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_E_VALUE);
        return CRYPT_RSA_ERR_E_VALUE;
    }
    return CRYPT_SUCCESS;
}

void KAE_RSA_FreePara(RSA_Para *para)
{
    if (para == NULL) {
        return;
    }

    BN_Destroy(para->e);
    BN_Destroy(para->p);
    BN_Destroy(para->q);

    free(para);
    para = NULL;
}

RSA_Para *KAE_RSA_NewPara(const BSL_Param *para)
{
    const uint8_t *e = NULL;
    uint32_t eLen = 0;
    int32_t ret = GetRsaParam(para, CRYPT_PARAM_RSA_E, &e, &eLen);
    if (ret != CRYPT_SUCCESS) {
        return NULL;
    }
    uint32_t bits = 0;
    ret = GetRsaBits(para, &bits);
    if (ret != CRYPT_SUCCESS) {
        return NULL;
    }
    ret = ValidateRsaParams(eLen, bits);
    if (ret != CRYPT_SUCCESS) {
        return NULL;
    }
    RSA_Para *retPara = (RSA_Para *)malloc(sizeof(RSA_Para));
    if (retPara == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }
    retPara->bits = bits;
    retPara->e = BN_Create(bits);
    retPara->p = BN_Create(bits);
    retPara->q = BN_Create(bits);

    if (retPara->e == NULL || retPara->p == NULL || retPara->q == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        goto ERR;
    }
    ret = BN_Bin2Bn(retPara->e, e, eLen);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    if (BN_BITS_TO_BYTES(bits) > RSA_SMALL_MODULUS_BYTES && BN_Bytes(retPara->e) > RSA_MAX_PUBEXP_BYTES) {
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_KEY_BITS);
        goto ERR;
    }
    return retPara;
ERR:
    KAE_RSA_FreePara(retPara);
    retPara = NULL;
    return NULL;
}

void KAE_RSA_FreePrvKey(KAE_RSA_PrvKey *prvKey)
{
    if (prvKey == NULL) {
        return;
    }
    BN_Destroy(prvKey->n);
    BN_Destroy(prvKey->d);
    BN_Destroy(prvKey->p);
    BN_Destroy(prvKey->q);
    BN_Destroy(prvKey->e);
    BN_Destroy(prvKey->dP);
    BN_Destroy(prvKey->dQ);
    BN_Destroy(prvKey->qInv);
    BSL_SAL_FREE(prvKey);
}

KAE_RSA_PrvKey *KAE_RSA_NewPrvKey(uint32_t bits)
{
    KAE_RSA_PrvKey *priKey = (KAE_RSA_PrvKey *)malloc(sizeof(KAE_RSA_PrvKey));
    if (priKey == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }
    priKey->n = BN_Create(bits);
    priKey->d = BN_Create(bits);
    priKey->p = BN_Create(bits);
    priKey->q = BN_Create(bits);
    priKey->e = BN_Create(bits);
    priKey->dP = BN_Create(bits);
    priKey->dQ = BN_Create(bits);
    priKey->qInv = BN_Create(bits);
    bool creatFailed = (priKey->n == NULL || priKey->d == NULL || priKey->e == NULL || priKey->p == NULL ||
                        priKey->q == NULL || priKey->dP == NULL || priKey->dQ == NULL || priKey->qInv == NULL);
    if (creatFailed) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        KAE_RSA_FreePrvKey(priKey);
        priKey = NULL;
    }
    return priKey;
}

void KAE_RSA_FreePubKey(KAE_RSA_PubKey *pubKey)
{
    if (pubKey == NULL) {
        return;
    }
    BN_Destroy(pubKey->n);
    BN_Destroy(pubKey->e);
    BN_MontDestroy(pubKey->mont);
    BSL_SAL_FREE(pubKey);
}

KAE_RSA_PubKey *KAE_RSA_NewPubKey(uint32_t bits)
{
    KAE_RSA_PubKey *pubKey = (KAE_RSA_PubKey *)malloc(sizeof(KAE_RSA_PubKey));
    if (pubKey == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }
    pubKey->n = BN_Create(bits);
    pubKey->e = BN_Create(bits);
    pubKey->mont = NULL;
    if (pubKey->n == NULL || pubKey->e == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        KAE_RSA_FreePubKey(pubKey);
        pubKey = NULL;
    }
    return pubKey;
}

int32_t GetRandomX(void *libCtx, BN_BigNum *X, uint32_t nlen, bool isP)
{
    /*
     *  The FIPS 185-5 Appendix B.9 required √2(2 ^(nlen/2 - 1)) <= x <= ((2 ^(nlen/2) - 1))
     *  hence we can limit it as follows:
     *      √2 ~= 1.41421 < 1.5 -->
     *          √2(2 ^(nlen/2 - 1)) < 1.5 * (2 ^(nlen/2 - 1))
     *          next, we need to prove 1.5 * (2 ^(nlen/2 - 1)) <= ((2 ^(nlen/2) - 1))
     *              --> let x = 2 ^(nlen/2),  1.5 * (x/2) ≤ x - 1
     *              --> (3/4) x ≤ x - 1
     *              --> x >= 4, obviously correct.
     *  And, 1.5 * 2 ^(nlen/2 - 1) = 2 ^ (nlen/2 - 1) + 2 ^ (nlen/2 - 2);
     *  If we follow these steps to construct the bigNum:
     *      i. Randomly generate a random number, the most significant bit is (nlen / 2).
     *      ii. Set the (nlen/2 - 1) bits.
     *  We can obtain the x, satisfied [ 1.5 * 2 ^(nlen/2 - 1), ((2 ^(nlen/2) - 1) ].
     */
    if ((nlen % 2) == 0) {
        return BN_RandEx(libCtx, X, nlen >> 1, BN_RAND_TOP_TWOBIT, BN_RAND_BOTTOM_NOBIT);
    }
    /*
     * Meanwhile, if nlen is odd, We need to consider p, q separately.
     */
    if (isP) {
        /*
         *  left : √2(2 ^(nlen/2 - 1)) < 2 ^ ⌊ (nlen / 2) ⌋
         *  right: if nlen is odd, 2 ^ (nlen/2) - 1 == 2 ^ ( ⌊ (nlen)/2 ⌋ + 1/2) - 1 == √2 * 2 ^ (⌊ (nlen)/2 ⌋) - 1
         *  if we want left <= right:
         *         2 ^ ⌊ (nlen / 2) ⌋ < √2 * 2 ^ (⌊ (nlen)/2 ⌋) - 1
         *    -->  2 ^ ⌊ (nlen / 2) ⌋ < 1.4 * 2 ^ (⌊ (nlen)/2 ⌋) - 1
         *    -->  1 < 0.4 * 2 ^ (⌊ (nlen)/2 ⌋)
         *    -->  nlen >= 3, obviously correct.
         *  hence, We can obtain the x, set the (nlen)/2 + 1 bits.
         */
        return BN_RandEx(libCtx, X, (nlen + 1) >> 1, BN_RAND_TOP_ONEBIT, BN_RAND_BOTTOM_NOBIT);
    }
    return BN_RandEx(libCtx, X, nlen >> 1, BN_RAND_TOP_TWOBIT, BN_RAND_BOTTOM_NOBIT);
}

/*
 * Ref: FIPS 186-5: Table A.1
 * Get the maximum lengths of p1, p2, q1, and q2.
 */
uint32_t GetAuxiliaryPrimeBitLen(uint32_t nlen)
{
    if (nlen <= 3071) {
        return 141;
    } else if (nlen <= 4095) {
        return 171;
    } else {
        return 201;
    }
}

/*
 * Ref: FIPS 186-5: Table A.1
 * Get the maximum lengths of p, q.
 */
uint32_t GetProbableNoLimitedBitLen(uint32_t nlen)
{
    if (nlen <= 3071) {
        return 1007;
    } else if (nlen <= 4095) {
        return 1518;
    } else {
        return 2030;
    }
}

/*
 * Ref: FIPS 186-5: Table B.1
 * Get minimum number of rounds of M-R testing when generating auxiliary primes.
 */
uint32_t GetAuxPrimeMillerCheckTimes(uint32_t auxBits)
{
    if (auxBits <= 170) {
        return 38; // Error probability = 2 ^ (-112)
    } else if (auxBits <= 200) {
        return 41; // Error probability = 2 ^ (-128)
    } else {
        return 44; // Error probability = 2 ^ (-144)
    }
}

/*
 * Ref: FIPS 186-5: Table B.1
 * Get minimum number of rounds of M-R testing when generating probable primes.
 */
uint32_t GetProbPrimeMillerCheckTimes(uint32_t proBits)
{
    if (proBits < 1536) {
        return 5;
    }
    return 4;
}

int32_t GenAuxPrime(BN_BigNum *Xp, uint32_t auxBits, BN_Optimizer *opt, bool isSeed)
{
    int32_t ret = CRYPT_SUCCESS;
    if (!isSeed) {
        ret = BN_RandEx(BN_OptimizerGetLibCtx(opt), Xp, auxBits, BN_RAND_TOP_ONEBIT, BN_RAND_BOTTOM_ONEBIT);
        if (ret != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(ret);
            return ret;
        }
    }
    uint32_t auxPrimeCheck = GetAuxPrimeMillerCheckTimes(auxBits);
    do {
        ret = BN_PrimeCheck(Xp, auxPrimeCheck, opt, NULL);
        if (ret == CRYPT_SUCCESS) {
            return ret;
        }
        if (ret != CRYPT_BN_NOR_CHECK_PRIME) {
            BSL_ERR_PUSH_ERROR(ret);
            return ret;
        }
        ret = BN_AddLimb(Xp, Xp, 2); // Try with odd numbers every time.
        if (ret != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(ret);
            return ret;
        }
    } while (true);
}

/*
 * Ref: FIPS 186-5 B.9 Compute a Probable Prime Factor Based on Auxiliary Primes.
 * The standard specifies that the length of two small primes should meet
 *                 len(r1) + len(r2) ≤ (nlen/2) – log2(nlen/2) – 7
 * If nlen = 1024, r1, r2 is obtained by search from 141 bits data, the above inequality is still satisfied.
 * Hence, it's a only performance consideration for us to use this standard for 1024-bit rsa key-Gen.
 */
int32_t GenPrimeWithAuxiliaryPrime(uint32_t auxBits, uint32_t proBits, BN_BigNum *Xp, BN_BigNum *Xp0,
    BN_BigNum *Xp1, BN_BigNum *Xp2, BN_BigNum *p, const CRYPT_RSA_Para *para, bool isP, BN_Optimizer *opt)
{
    BN_BigNum *r1;
    BN_BigNum *r2;
    uint32_t auxRoom = BITS_TO_BN_UNIT(auxBits);
    int32_t ret = OptimizerStart(opt); // use the optimizer
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }
    uint32_t probPrimeCheck = GetProbPrimeMillerCheckTimes(proBits);

    r1 = (Xp1 != NULL) ? Xp1 : OptimizerGetBn(opt, auxRoom);
    r2 = (Xp2 != NULL) ? Xp2 : OptimizerGetBn(opt, auxRoom);

    BN_BigNum *r1Double = OptimizerGetBn(opt, auxRoom);
    BN_BigNum *primeCheck = OptimizerGetBn(opt, auxRoom);
    BN_BigNum *r2Inv = OptimizerGetBn(opt, auxRoom);
    BN_BigNum *r1DoubleInv = OptimizerGetBn(opt, auxRoom);
    BN_BigNum *R = OptimizerGetBn(opt, auxRoom);
    BN_BigNum *pMinusOne = OptimizerGetBn(opt, BITS_TO_BN_UNIT(proBits));
    uint32_t bits = isP ? (para->bits + 1) >> 1 : (para->bits >> 1); // Avoid the bit is odd.
    uint32_t iterRound = 20 * bits; // Step 9 specifies that the iteration round is 20 * (nlen/2);
    if (r1 == NULL || r2 == NULL || r1Double == NULL || primeCheck == NULL || r2Inv == NULL ||
        r1DoubleInv == NULL || R == NULL || pMinusOne == NULL) {
        ret = CRYPT_BN_OPTIMIZER_GET_FAIL;
        OptimizerEnd(opt);
        return ret;
    }

    // Choose auxiliary prime r1, either from seed or generate randomly
    ret = GenAuxPrime(r1, auxBits, opt, (Xp1 != NULL));
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        OptimizerEnd(opt);
        return ret;
    }
    GOTO_ERR_IF(GenAuxPrime(r2, auxBits, opt, (Xp2 != NULL)), ret);
    GOTO_ERR_IF(BN_Lshift(r1Double, r1, 1), ret);
    // Step 1: check 2r1, r2 are coprime.
    GOTO_ERR_IF(BN_Gcd(primeCheck, r1Double, r2, opt), ret);
    if (!BN_IsOne(primeCheck)) {
        ret = CRYPT_RSA_NOR_KEYGEN_FAIL;
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_NOR_KEYGEN_FAIL);
        goto ERR;
    }
    // Step 2: cal R = (r2^-1 mod 2r1) * r2 - ((2 * r1)^-1 mod r2) * (2 * r1)
    GOTO_ERR_IF(BN_ModInv(r2Inv, r2, r1Double, opt), ret); // (r2^-1 mod 2r1) * r2
    GOTO_ERR_IF(BN_Mul(r2Inv, r2, r2Inv, opt), ret);
    // ((2 * r1)^-1 mod r2) * (2 * r1)
    GOTO_ERR_IF(BN_ModInv(r1DoubleInv, r1Double, r2, opt), ret);
    GOTO_ERR_IF(BN_Mul(r1DoubleInv, r1Double, r1DoubleInv, opt), ret);
    // get R.
    GOTO_ERR_IF(BN_Sub(R, r2Inv, r1DoubleInv), ret);
    do {
        // Step 3: get x via seed xp/xq or random
        if (Xp0 == NULL) {
            GOTO_ERR_IF(GetRandomX(BN_OptimizerGetLibCtx(opt), Xp, para->bits, isP), ret);
        }

        // Step 4: Y = X + ((R – X) mod 2r1r2
        GOTO_ERR_IF(BN_Mul(r1, r1Double, r2, opt), ret); // 2r1r2
        GOTO_ERR_IF(BN_ModSub(R, R, Xp, r1, opt), ret);
        GOTO_ERR_IF(BN_Add(p, Xp, R), ret);
        uint32_t i = 0;
        for (; i < iterRound; i++) {
            // Step 6: Check p ≥ 2 ^ (nlen/2)
            if (BN_Bits(p) > bits) {
                break;
            }
            // Step 7: Check the p - 1 and e are corprime.
            GOTO_ERR_IF(BN_SubLimb(pMinusOne, p, 1), ret);
            GOTO_ERR_IF(BN_Gcd(pMinusOne, pMinusOne, para->e, opt), ret);
            if (BN_IsOne(pMinusOne)) {
                // Step 7.1: Check the primality of p.
                ret = BN_PrimeCheck(p, probPrimeCheck, opt, NULL);
                if (ret == CRYPT_SUCCESS) { // We find a primes successfully.
                    goto ERR;
                }
                if (ret != CRYPT_BN_NOR_CHECK_PRIME) { // Another exception has occurred.
                    BSL_ERR_PUSH_ERROR(ret);
                    goto ERR;
                }
            }
            // Step 10: Update p.
            GOTO_ERR_IF(BN_Add(p, p, r1), ret);
        }
        // Step 9: check i ≥ 20 * (nlen/2).
        if (i == iterRound) {
            ret = CRYPT_RSA_NOR_KEYGEN_FAIL;
            BSL_ERR_PUSH_ERROR(ret);
            goto ERR;
        }
    } while (true);
ERR:
    if (Xp1 == NULL) {
        BN_Zeroize(r1);
    }
    if (Xp2 == NULL) {
        BN_Zeroize(r2);
    }
    OptimizerEnd(opt);
    return ret;
}

int32_t GenPQBasedOnProbPrimes(const RSA_Para *para, KAE_RSA_PrvKey *priKey, BN_Optimizer *opt)
{
    BN_BigNum *Xp = NULL, *Xq = NULL, *Xp0 = NULL, *Xp1 = NULL, *Xp2 = NULL, *Xq0 = NULL, *Xq1 = NULL, *Xq2 = NULL;
    uint32_t proBits = GetProbableNoLimitedBitLen(para->bits);
    uint32_t auxBits = GetAuxiliaryPrimeBitLen(para->bits);
    // Used in check |Xp – Xq| ≤ 2^(nlen/2) – 100 or |p – q| ≤ 2^(nlen/2) – 100.
    uint32_t secBits = ((para->bits + 1) >> 1) - 100;
    uint32_t proRoom = BITS_TO_BN_UNIT(proBits);
    int32_t ret = OptimizerStart(opt); // use the optimizer
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }

    Xp = (Xp0 != NULL) ? Xp0 : OptimizerGetBn(opt, proRoom);
    Xq = (Xq0 != NULL) ? Xq0 : OptimizerGetBn(opt, proRoom);
    if (Xp == NULL || Xq == NULL) {
        ret = CRYPT_BN_OPTIMIZER_GET_FAIL;
        BSL_ERR_PUSH_ERROR(ret);
        OptimizerEnd(opt);
        return ret;
    }

    // Step 4: get p
    ret = GenPrimeWithAuxiliaryPrime(auxBits, proBits, Xp, Xp0, Xp1, Xp2, priKey->p, para, true, opt);
    if (ret != CRYPT_SUCCESS) {
        BN_Zeroize(Xp);
        BSL_ERR_PUSH_ERROR(ret);
        OptimizerEnd(opt);
        return ret;
    }
    /*
     * If |Xp – Xq| ≤ 2 ^ (2nlen/2 – 100) or |p – q| ≤ 2 ^ (2nlen/2 – 100), need to try again.
     * We think there can ever be repeated many times here unless the 'random' is stuck.
     * For example, nlen = 2048 and |Xp – Xq| ≤ 2 ^ (1024 – 100), it means that the most significant
     * 99 bits of our Xq and Xp randomly generated are all identical. It's a low-probability event.
     */
    do {
        // Step 5: get q
        ret = GenPrimeWithAuxiliaryPrime(auxBits, proBits, Xq, Xq0, Xq1, Xq2, priKey->q, para, false, opt);
        if (ret != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(ret);
            goto ERR;
        }
        // Step 6: Check (|Xp – Xq| ≤ 2^(nlen/2) – 100) and (|p – q| ≤ 2^(nlen/2) – 100)
        ret = BN_Sub(Xq, Xp, Xq); // Xq dont needs anymore, but Xp may be used.
        if (ret != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(ret);
            goto ERR;
        }
        // |Xp – Xq| ≤ 2 ^ (2nlen/2 – 100) -> BN_Bits(Xp) <= secBits + 1 -> BN_Bits(Xp) < secBits
        if (BN_Bits(Xq) < secBits) {
            if (Xq0 != NULL && Xq1 != NULL && Xq2 != NULL) {
                ret = CRYPT_RSA_NOR_KEYGEN_FAIL;
                BSL_ERR_PUSH_ERROR(ret);
                goto ERR;
            }
            continue;
        }
        ret = BN_Sub(Xq, priKey->p, priKey->q);
        if (ret != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(ret);
            goto ERR;
        }
        // |p – q| ≤ 2 ^ (2nlen/2 – 100)
        if (BN_Bits(Xq) < secBits) {
            if (Xq0 != NULL && Xq1 != NULL && Xq2 != NULL) {
                ret = CRYPT_RSA_NOR_KEYGEN_FAIL;
                BSL_ERR_PUSH_ERROR(ret);
                goto ERR;
            }
            continue;
        }
        break;
    } while (true);
ERR:
    if (Xp0 == NULL) {
        BN_Zeroize(Xp);
    }
    if (Xq0 == NULL) {
        BN_Zeroize(Xq);
    }
    OptimizerEnd(opt);
    return ret;
}

int32_t KAE_RSA_Gen(KAE_RSA_Ctx *ctx)
{
    int32_t ret = CRYPT_MEM_ALLOC_FAIL;
    BN_Optimizer *optimizer = NULL;

    optimizer = BN_OptimizerCreate();
    if (optimizer == NULL || ctx->kPrvKey == NULL || ctx->kPubKey == NULL) {
        ret = CRYPT_MEM_ALLOC_FAIL;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    ret = GenPQBasedOnProbPrimes(ctx->para, ctx->kPrvKey, optimizer);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
ERR:
    BN_OptimizerDestroy(optimizer);
    return ret;
}

uint32_t KAE_RSA_GetBits(const KAE_RSA_Ctx *ctx)
{
    if (ctx == NULL) {
        return 0;
    }
    if (ctx->para != NULL) {
        return ctx->para->bits;
    }
    if (ctx->kPrvKey != NULL) {
        return BN_Bits(ctx->kPrvKey->n);
    }
    if (ctx->kPubKey != NULL) {
        return BN_Bits(ctx->kPubKey->n);
    }
    return 0;
}

uint32_t KAE_RSA_GetSignLen(const KAE_RSA_Ctx *ctx)
{
    return BN_BITS_TO_BYTES(KAE_RSA_GetBits(ctx));
}

int32_t KAE_RSA_GetSecBits(const KAE_RSA_Ctx *ctx)
{
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return 0;
    }
    int32_t bits = (int32_t)KAE_RSA_GetBits(ctx);
    return BN_SecBits(bits, -1);
}

int32_t KAE_RSA_SetPara(KAE_RSA_Ctx *ctx, const BSL_Param *para)
{
    if (ctx == NULL || para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    RSA_Para *rsaPara = KAE_RSA_NewPara(para);
    if (rsaPara == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_EAL_ERR_NEW_PARA_FAIL);
        return CRYPT_EAL_ERR_NEW_PARA_FAIL;
    }
    int32_t ret = IsRSASetParaVaild(ctx, rsaPara);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        KAE_RSA_FreePara(rsaPara);
        rsaPara = NULL;
        return ret;
    }

    (void)memset_s(&(ctx->pad), sizeof(RSAPad), 0, sizeof(RSAPad));
    KAE_RSA_FreePara(ctx->para);
    ctx->para = NULL;
    KAE_RSA_FreePrvKey(ctx->kPrvKey);
    ctx->kPrvKey = NULL;
    KAE_RSA_FreePubKey(ctx->kPubKey);
    ctx->kPubKey = NULL;
    ctx->para = rsaPara;
    ctx->setup.key_bits = rsaPara->bits;

    return CRYPT_SUCCESS;
}

int32_t KAE_RSA_SetPub(KAE_RSA_Ctx *ctx, const BSL_Param *para)
{
    BSL_Param *ncPara = (BSL_Param *)para;
    BSL_Param *n = BSL_PARAM_FindParam(ncPara, CRYPT_PARAM_RSA_N);
    BSL_Param *e = BSL_PARAM_FindParam(ncPara, CRYPT_PARAM_RSA_E);

    if (ctx == NULL || para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    if (PARAMISNULL(n)) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    if (PARAMISNULL(e)) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    if (n->valueLen > RSA_MAX_MODULUS_LEN) {
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_KEY_BITS);
        return CRYPT_RSA_ERR_KEY_BITS;
    }
    if (e->valueLen > n->valueLen) {
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_KEY_BITS);
        return CRYPT_RSA_ERR_KEY_BITS;
    }

    uint32_t bnBits;
    KAE_RSA_PubKey *newPub = NULL;
    (void)memset_s(&(ctx->pad), sizeof(RSAPad), 0, sizeof(RSAPad));

    newPub = KAE_RSA_NewPubKey(n->valueLen * 8);
    if (newPub == NULL) {
        return CRYPT_MEM_ALLOC_FAIL;
    }
    int32_t ret;
    GOTO_ERR_IF(BN_Bin2Bn(newPub->n, n->value, n->valueLen), ret);
    bnBits = BN_Bits(newPub->n);
    if (bnBits > RSA_MAX_MODULUS_BITS || bnBits < RSA_MIN_MODULUS_BITS) {
        ret = CRYPT_RSA_ERR_KEY_BITS;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    GOTO_ERR_IF(BN_Bin2Bn(newPub->e, e->value, e->valueLen), ret);
    if (n->valueLen > RSA_SMALL_MODULUS_BYTES && BN_Bytes(newPub->e) > RSA_MAX_PUBEXP_BYTES) {
        ret = CRYPT_RSA_ERR_KEY_BITS;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }

    if (BN_Cmp(newPub->n, newPub->e) <= 0 || BN_IsZero(newPub->e) || BN_IsOne(newPub->e)) {
        ret = CRYPT_RSA_ERR_INPUT_VALUE;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }

    newPub->mont = BN_MontCreate(newPub->n);
    if (newPub->mont == NULL) {
        ret = CRYPT_MEM_ALLOC_FAIL;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }

    KAE_RSA_FreePubKey(ctx->kPubKey);
    ctx->kPubKey = newPub;
    ctx->setup.key_bits = n->valueLen * 8;
    return ret;
ERR:
    KAE_RSA_FreePubKey(newPub);
    newPub = NULL;
    return ret;
}

static int32_t RsaPrvKeyCalcND(KAE_RSA_Ctx *ctx, BN_BigNum *pMinusOne, BN_BigNum *qMinusOne, BN_Optimizer *optimizer)
{
    int32_t ret;
    if (ctx->para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    KAE_RSA_PrvKey *prvKey = ctx->kPrvKey;
    BN_BigNum *l = BN_Create(ctx->para->bits);
    BN_BigNum *u = BN_Create(ctx->para->bits);
    if (l == NULL || u == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        ret = CRYPT_MEM_ALLOC_FAIL;
        goto EXIT;
    }
    ret = BN_Mul(prvKey->n, prvKey->p, prvKey->q, optimizer);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }
    ret = BN_Mul(l, pMinusOne, qMinusOne, optimizer);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }
    ret = BN_Gcd(u, pMinusOne, qMinusOne, optimizer);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }
    ret = BN_Div(l, NULL, l, u, optimizer);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }
    ret = BN_ModInv(prvKey->d, ctx->para->e, l, optimizer);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
    }
EXIT:
    BN_Destroy(l);
    BN_Destroy(u);
    return ret;
}

// p, q [ => n, d]  => dP dQ qInv
// ctx->para may be NULL when setting key
int32_t KAE_RSA_CalcPrvKey(KAE_RSA_Ctx *ctx, BN_Optimizer *optimizer)
{
    int32_t ret;
    KAE_RSA_PrvKey *prvKey = ctx->kPrvKey;
    BN_BigNum *pMinusOne = BN_Create(BN_Bits(prvKey->p));
    BN_BigNum *qMinusOne = BN_Create(BN_Bits(prvKey->q));
    if (pMinusOne == NULL || qMinusOne == NULL) {
        ret = CRYPT_MEM_ALLOC_FAIL;
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }
    ret = BN_SubLimb(pMinusOne, prvKey->p, 1);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }
    ret = BN_SubLimb(qMinusOne, prvKey->q, 1);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }
    if (BN_IsZero(prvKey->n)) { // when generating key
        ret = RsaPrvKeyCalcND(ctx, pMinusOne, qMinusOne, optimizer);
        if (ret != CRYPT_SUCCESS) {
            goto EXIT;
        }
    }
    ret = BN_ModInv(prvKey->qInv, prvKey->q, prvKey->p, optimizer);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }
    ret = BN_Div(NULL, prvKey->dP, prvKey->d, pMinusOne, optimizer);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }
    ret = BN_Div(NULL, prvKey->dQ, prvKey->d, qMinusOne, optimizer);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
    }
EXIT:
    BN_Destroy(pMinusOne);
    BN_Destroy(qMinusOne);
    return ret;
}

int32_t KAE_RSA_SetPrv(KAE_RSA_Ctx *ctx, const BSL_Param *para)
{
    if (ctx == NULL || para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    BSL_Param *ncPara = (BSL_Param *)para;
    BSL_Param *n = BSL_PARAM_FindParam(ncPara, CRYPT_PARAM_RSA_N);
    BSL_Param *d = BSL_PARAM_FindParam(ncPara, CRYPT_PARAM_RSA_D);
    BSL_Param *e = BSL_PARAM_FindParam(ncPara, CRYPT_PARAM_RSA_E);
    BSL_Param *p = BSL_PARAM_FindParam(ncPara, CRYPT_PARAM_RSA_P);
    BSL_Param *q = BSL_PARAM_FindParam(ncPara, CRYPT_PARAM_RSA_Q);
    BSL_Param *dP = BSL_PARAM_FindParam(ncPara, CRYPT_PARAM_RSA_DP);
    BSL_Param *dQ = BSL_PARAM_FindParam(ncPara, CRYPT_PARAM_RSA_DQ);
    BSL_Param *qInv = BSL_PARAM_FindParam(ncPara, CRYPT_PARAM_RSA_QINV);

    if (PARAMISNULL(n) || n->valueLen == 0 || PARAMISNULL(d)) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    if (n->valueLen > RSA_MAX_MODULUS_LEN) {
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_KEY_BITS);
        return CRYPT_RSA_ERR_KEY_BITS;
    }
    // prv->p\q and prv->dP\dQ\qInv must be both empty or not.
    // If prv->p is empty, prv->dP must be empty.
    if ((PARAMISNULL(p) != PARAMISNULL(q)) || (PARAMISNULL(p) && !PARAMISNULL(dP))) {
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_NO_KEY_INFO);
        return CRYPT_RSA_NO_KEY_INFO;
    }
    if ((PARAMISNULL(dP) || PARAMISNULL(dQ) || PARAMISNULL(qInv)) &&
        (!PARAMISNULL(dP) || !PARAMISNULL(dQ) || !PARAMISNULL(qInv))) {
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_NO_KEY_INFO);
        return CRYPT_RSA_NO_KEY_INFO;
    }

    uint32_t bnBytes = n->valueLen;
    if (d->valueLen > bnBytes || p->valueLen > bnBytes || q->valueLen > bnBytes) {
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_KEY_BITS);
        return CRYPT_RSA_ERR_KEY_BITS;
    }

    KAE_RSA_Ctx *newCtx = KAE_RSA_NewCtx();
    if (newCtx == NULL) {
        return CRYPT_MEM_ALLOC_FAIL;
    }

    int32_t ret;
    newCtx->kPrvKey = KAE_RSA_NewPrvKey(n->valueLen * 8); // Bit length is obtained by multiplying byte length by 8.
    if (newCtx->kPrvKey == NULL) {
        ret = CRYPT_MEM_ALLOC_FAIL;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }

    ret = BN_Bin2Bn(newCtx->kPrvKey->n, n->value, n->valueLen);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    uint32_t bnBits = BN_Bits(newCtx->kPrvKey->n);
    if (bnBits > RSA_MAX_MODULUS_BITS || bnBits < RSA_MIN_MODULUS_BITS) {
        ret = CRYPT_RSA_ERR_KEY_BITS;
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_KEY_BITS);
        goto ERR;
    }
    ret = BN_Bin2Bn(newCtx->kPrvKey->d, d->value, d->valueLen);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    // d cannot be 0 or 1. The mathematical logic of e and d is that
    // d and e are reciprocal in mod((p-1) * (q-1)); When d is 1, e and d must be 1. When d is 0, e doesn't exist.
    if (BN_IsZero(newCtx->kPrvKey->d) || BN_IsOne(newCtx->kPrvKey->d)) {
        ret = CRYPT_RSA_ERR_INPUT_VALUE;
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_INPUT_VALUE);
        goto ERR;
    }
    if (BN_Cmp(newCtx->kPrvKey->n, newCtx->kPrvKey->d) <= 0) {
        ret = CRYPT_RSA_ERR_INPUT_VALUE;
        BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_INPUT_VALUE);
        goto ERR;
    }
    if (!PARAMISNULL(e)) {
        ret = BN_Bin2Bn(newCtx->kPrvKey->e, e->value, e->valueLen);
        if (ret != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(ret);
            goto ERR;
        }
        if (BN_Cmp(newCtx->kPrvKey->n, newCtx->kPrvKey->e) <= 0) {
            ret = CRYPT_RSA_ERR_INPUT_VALUE;
            BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_INPUT_VALUE);
            goto ERR;
        }
    }
    if (!PARAMISNULL(p)) {
        GOTO_ERR_IF_EX(BN_Bin2Bn(newCtx->kPrvKey->p, p->value, p->valueLen), ret);
        GOTO_ERR_IF_EX(BN_Bin2Bn(newCtx->kPrvKey->q, q->value, q->valueLen), ret);
        if (BN_IsZero(newCtx->kPrvKey->p) == true || BN_IsZero(newCtx->kPrvKey->q) == true) {
            ret = CRYPT_RSA_ERR_INPUT_VALUE;
            BSL_ERR_PUSH_ERROR(CRYPT_RSA_ERR_INPUT_VALUE);
            goto ERR;
        }
        if (!PARAMISNULL(dP)) {
            GOTO_ERR_IF_EX(BN_Bin2Bn(newCtx->kPrvKey->dP, dP->value, dP->valueLen), ret);
            GOTO_ERR_IF_EX(BN_Bin2Bn(newCtx->kPrvKey->dQ, dQ->value, dQ->valueLen), ret);
            GOTO_ERR_IF_EX(BN_Bin2Bn(newCtx->kPrvKey->qInv, qInv->value, qInv->valueLen), ret);
        }
    }

    if (!PARAMISNULL(p) && PARAMISNULL(dP)) {
        BN_Optimizer *optimizer = BN_OptimizerCreate();
        if (optimizer == NULL) {
            ret = CRYPT_MEM_ALLOC_FAIL;
            BSL_ERR_PUSH_ERROR(ret);
            goto ERR;
        }
        ret = KAE_RSA_CalcPrvKey(newCtx, optimizer);
        BN_OptimizerDestroy(optimizer);
        if (ret != CRYPT_SUCCESS) {
            goto ERR;
        }
    }

    KAE_RSA_FreePrvKey(ctx->kPrvKey);
    ctx->kPrvKey = newCtx->kPrvKey;
    newCtx->kPrvKey = NULL;
    ctx->pad = newCtx->pad;
    ctx->setup.key_bits = n->valueLen * 8;
    ctx->setup.is_crt = !PARAMISNULL(p);

    KAE_RSA_Free(newCtx);
    return ret;
ERR:
    KAE_RSA_Free(newCtx);
    return ret;
}

int32_t KAE_RSA_GetPub(const KAE_RSA_Ctx *ctx, BSL_Param *para)
{
    if (ctx == NULL || ctx->kPubKey == NULL || para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    BSL_Param *e = BSL_PARAM_FindParam(para, CRYPT_PARAM_RSA_E);
    if (e == NULL || e->value == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    uint32_t eLen = e->valueLen;
    int32_t ret = BN_Bn2Bin(ctx->kPubKey->e, e->value, &eLen);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }

    BSL_Param *n = BSL_PARAM_FindParam(para, CRYPT_PARAM_RSA_N);
    if (n == NULL || n->value == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    uint32_t nLen = n->valueLen;
    ret = BN_Bn2Bin(ctx->kPubKey->n, n->value, &nLen);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }
    e->useLen = eLen;
    n->useLen = nLen;
    return CRYPT_SUCCESS;
}

int32_t KAE_RSA_GetPrv(const KAE_RSA_Ctx *ctx, BSL_Param *para)
{
    if (ctx == NULL || ctx->kPrvKey == NULL || para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    BSL_Param *n = BSL_PARAM_FindParam(para, CRYPT_PARAM_RSA_N);
    BSL_Param *d = BSL_PARAM_FindParam(para, CRYPT_PARAM_RSA_D);
    BSL_Param *e = BSL_PARAM_FindParam(para, CRYPT_PARAM_RSA_E);
    BSL_Param *p = BSL_PARAM_FindParam(para, CRYPT_PARAM_RSA_P);
    BSL_Param *q = BSL_PARAM_FindParam(para, CRYPT_PARAM_RSA_Q);
    BSL_Param *dP = BSL_PARAM_FindParam(para, CRYPT_PARAM_RSA_DP);
    BSL_Param *dQ = BSL_PARAM_FindParam(para, CRYPT_PARAM_RSA_DQ);
    BSL_Param *qInv = BSL_PARAM_FindParam(para, CRYPT_PARAM_RSA_QINV);

    if ((PARAMISNULL(p) != PARAMISNULL(q)) ||
        ((PARAMISNULL(dP) || PARAMISNULL(dQ) || PARAMISNULL(qInv)) &&
         (!PARAMISNULL(dP) || !PARAMISNULL(dQ) || !PARAMISNULL(qInv))) ||
        (PARAMISNULL(p) && !PARAMISNULL(dP))) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    int32_t ret;
    n->useLen = n->valueLen;
    GOTO_ERR_IF(BN_Bn2Bin(ctx->kPrvKey->n, n->value, &(n->useLen)), ret);
    d->useLen = d->valueLen;
    GOTO_ERR_IF(BN_Bn2Bin(ctx->kPrvKey->d, d->value, &(d->useLen)), ret);
    if (!PARAMISNULL(e)) {
        e->useLen = e->valueLen;
        GOTO_ERR_IF(BN_Bn2Bin(ctx->kPrvKey->e, e->value, &(e->useLen)), ret);
    }
    if (!PARAMISNULL(p)) {
        p->useLen = p->valueLen;
        GOTO_ERR_IF(BN_Bn2Bin(ctx->kPrvKey->p, p->value, &(p->useLen)), ret);

        q->useLen = q->valueLen;
        GOTO_ERR_IF(BN_Bn2Bin(ctx->kPrvKey->q, q->value, &(q->useLen)), ret);
    }
    if (!PARAMISNULL(dQ)) {
        dQ->useLen = dQ->valueLen;
        GOTO_ERR_IF(BN_Bn2Bin(ctx->kPrvKey->dQ, dQ->value, &(dQ->useLen)), ret);
        dP->useLen = dP->valueLen;
        GOTO_ERR_IF(BN_Bn2Bin(ctx->kPrvKey->dP, dP->value, &(dP->useLen)), ret);
        qInv->useLen = qInv->valueLen;
        GOTO_ERR_IF(BN_Bn2Bin(ctx->kPrvKey->qInv, qInv->value, &(qInv->useLen)), ret);
    }
ERR:
    return ret;
}